package dilithium

import (
	"encoding/hex"
	"fmt"
	"test/gincircl/common"

	"github.com/cloudflare/circl/sign"
	"github.com/cloudflare/circl/sign/dilithium/mode2"
	"github.com/cloudflare/circl/sign/dilithium/mode3"
	"github.com/cloudflare/circl/sign/dilithium/mode5"
)

type DilithiumrWrapper struct {
	scheme sign.Scheme
}

// 创建新的 Kyber 实例
func NewDilithiumr(variant string) (*DilithiumrWrapper, error) {
	switch variant {
	case "dilithium2":
		return &DilithiumrWrapper{scheme: mode2.Scheme()}, nil
	case "dilithium3":
		return &DilithiumrWrapper{scheme: mode3.Scheme()}, nil
	case "dilithium5":
		return &DilithiumrWrapper{scheme: mode5.Scheme()}, nil
	default:
		return nil, fmt.Errorf("不支持的 Kyber 变体: %s", variant)
	}
}

// GenerateKey 生成 Dilithium 密钥对
func (k *DilithiumrWrapper) GenerateKey() (string, string, error) {

	// 生成公私钥对
	pk, sk, err := k.scheme.GenerateKey()
	if err != nil {
		common.Logger.Error("生成密钥对失败:" + err.Error())
		return "", "", err
	}
	pkBytes, err := pk.MarshalBinary()
	if err != nil {
		common.Logger.Error("公钥序列化失败:" + err.Error())
		return "", "", err
	}

	skBytes, err := sk.MarshalBinary()
	if err != nil {
		common.Logger.Error("私钥序列化失败:" + err.Error())
		return "", "", err
	}
	return hex.EncodeToString(pkBytes), hex.EncodeToString(skBytes), nil
}

// Sign 使用私钥对消息进行签名
func (k *DilithiumrWrapper) Sign(privateKey string, message string, opts *sign.SignatureOpts) (string, error) {
	skBytes, err := hex.DecodeString(privateKey)
	if err != nil {
		common.Logger.Error("解码私钥失败:" + err.Error())
		return "", fmt.Errorf("解码私钥失败: %v", err)
	}
	priKey, err := k.scheme.UnmarshalBinaryPrivateKey(skBytes)
	if err != nil {
		common.Logger.Error("私钥反序列化失败:" + err.Error())
		return "", fmt.Errorf("私钥反序列化失败: %v", err)
	}
	signText := k.scheme.Sign(priKey, []byte(message), opts)
	// 签名
	return hex.EncodeToString(signText), nil
}

// Verify 使用公钥验证签名
func (k *DilithiumrWrapper) Verify(publicKey string, message, signature string, opts *sign.SignatureOpts) (bool, error) {

	pkBytes, err := hex.DecodeString(publicKey)
	if err != nil {
		common.Logger.Error("解码公钥失败:" + err.Error())
		return false, fmt.Errorf("解码公钥失败: %v", err)
	}
	pubKey, err := k.scheme.UnmarshalBinaryPublicKey(pkBytes)
	if err != nil {
		common.Logger.Error("公钥反序列化失败:" + err.Error())
		return false, fmt.Errorf("公钥反序列化失败: %v", err)
	}
	signatureBytes, err := hex.DecodeString(signature)
	if err != nil {
		common.Logger.Error("解码签名失败:" + err.Error())
		return false, fmt.Errorf("解码签名失败: %v", err)
	}
	match := k.scheme.Verify(pubKey, []byte(message), signatureBytes, opts)
	// 验证签名
	return match, nil
}
